from torch_geometric.nn import SAGPooling, GCNConv,GATv2Conv,GatedGraphConv
from torch_geometric.nn import global_mean_pool as gap, global_max_pool as gmp
from torch_geometric.nn import MessagePassing
from torch_geometric.nn.conv.gcn_conv import GCNConv

import torch.nn as nn
import torch
import torch.nn.functional as F
from typing import List, Tuple

if __name__=="__main__":

    from embedding_layer import EmbeddingLayer
    from mlp_layer import MLPLayer
    from inter_conv import InterConv
    
else:
    from .embedding_layer import EmbeddingLayer
    from .mlp_layer import MLPLayer
    from .inter_conv import InterConv

class Tensor_GCN(MessagePassing):
    def __init__(self,
                 num_edge_types,
                 in_features,
                 out_features,
                 embedding_out_features,
                 embedding_num_classes,
                 dropout=0,
                 max_node_per_graph=50,
                 add_self_loops=False,
                 bias=True,
                 aggr="mean",
                 device="cpu"):
        super(Tensor_GCN, self).__init__(aggr=aggr)
        # params set
        self.num_edge_types = num_edge_types
        self.device = device
        self.dropout = dropout
        self.in_features=in_features
        self.out_features=out_features
        self.max_node_per_graph=max_node_per_graph
        self.embeding_out_features=embedding_out_features
        # 先对值进行embedding
        self.value_embeddingLayer = EmbeddingLayer(embedding_num_classes,
                                                   in_features,
                                                   embedding_out_features,
                                                   device=device)



        #self.rnns=[torch.nn.GRU(embedding_out_features,out_features,batch_first=True,device=device) for _ in range(2)]
        
        #self.conv1_list = nn.ModuleList([GCNConv(out_features, out_features, add_self_loops=add_self_loops) for _ in range(self.num_edge_types)] )
        #self.interConv = InterConv(out_features, out_features, num_edge_types=self.num_edge_types)
        #self.thirdNorm = torch.nn.InstanceNorm2d(out_features)  # 这个就是自己要找的
        
        
        # self.gat_time_step=[3,1,3,1]
        
        # self.gat_layers=[
        #         GATv2Conv(out_features, out_features//4, add_self_loops=False,heads=4)
        #         for _ in range(4)
        #     ]
        
        self.ggnns=[GatedGraphConv(out_features,8).to(self.device) for _ in range(num_edge_types)]
        
        self.conv1 = GATv2Conv(out_features, out_features//8,add_self_loops=False, heads=8, concat=True)
        

        #self.transformer=torch.nn.TransformerEncoder(torch.nn.TransformerEncoderLayer(out_features,4,batch_first=True),3)
    # def forward(self, x, edge_list: List[torch.tensor], **kwargs):
    #     # 8 层的tensor 卷积。效果也还不错，但是收敛速度没有添加ggnn的快。
    #     x_embedding = self.value_embeddingLayer(x)
    #     # Tensor GGNN 
    #     last_node_states = x_embedding
        
    #     loop_edge_list = self.matrix_loop_new(edge_list) # 4V, 4V
    #     for _ in range(8):

    #         #ggnn_out = x_embedding
    #         # tensor GCN:
    #         # 直接8层
    #         cur_x = torch.cat([last_node_states for _ in range(self.num_edge_types)], dim=0)  # 4V, D
    #         out = self.conv1(cur_x, loop_edge_list) # 4V, D
    #         out = out.view(self.num_edge_types, x_embedding.shape[0], out.shape[-1])  # 4, V, D
    #         out = torch.sum(out, dim=0)  # V, D
    #         out = F.relu(out)
    #         last_node_states = out


    #     return last_node_states
    
    def forward(self, x, edge_list: List[torch.tensor], **kwargs):
        # 这个是比较不错的代码的一个备份
        #print(x.size())
        def graph_compute(_model,_input,_edges,_batch_size):
            graph_output_list=[]
            for i in range(_batch_size):
                graph_output_list.append(_model(_input[i],_edges)) 
            
            return torch.stack(graph_output_list)
                
        #print(sum([e.size(1) for e in edge_list]))
        
        #print(batch_size)
        
        x_reshape=x.view(-1,self.in_features)
        x_embedding = self.value_embeddingLayer(x_reshape)
        
        
        #print(x_embedding.size())
        
        
        #嵌入邻域信息和语法信息
        ggnn_out_list=[]
        for i,ggnn in enumerate(self.ggnns):
            ggnn_out_list.append(ggnn(x_embedding,edge_list[i]))
            
        ggnn_out=sum(ggnn_out_list)
        
        # # 使用RNN完成单词的基于上下文的特征表示 status.size()=(batch,seq_length,hidden_dim)
        # status=x_embedding
        # for i in range(len(self.rnns)):
        #     layer=self.rnns[i]
        #     status,_=layer(status)
        
        # 特征image化
        # status.size()=(batch,num_edge_types*seq_length,hidden_dim)
        # edges is a list,which element.size()=(batch,2,num_edge_types*edge_number)
            
        edges=self.matrix_loop_new(edge_list,x_reshape.size(0))
        
        ggnn_out = torch.cat([ggnn_out for _ in range(self.num_edge_types)], dim=0)  # 4V, D
        
        gat_out=self.conv1(ggnn_out,edges)
        gat_out=gat_out.view(self.num_edge_types,-1,self.max_node_per_graph,self.out_features)
        gat_out=torch.sum(gat_out,0)
        # 图注意力机制
        # status.size()=(batch,num_edge_types*seq_length,hidden_dim)
        # for i,time_step in enumerate(self.gat_time_step):
        #     layer=self.gat_layers[i].to(self.device)
        #     for _ in range(time_step):
        #         status_by_batch=[]
        #         for j in range(batch_size):
        #             status_by_batch.append(layer(status[j],edges[j]))
        #         status=torch.stack(status_by_batch)

        #gg_out=status.view(batch_size,self.num_edge_types,-1,self.out_features)
        #gg_out = torch.mean(gg_out, dim=1)  # V, D
        
        
        #transformer encoder
        
        return torch.dropout(gat_out,self.dropout,self.training)


    def matrix_transfer(self, edge, i, j):
        # edge: [[i],[j]]
        edge_new = edge.detach().clone()
        edge_new[0]+=i
        edge_new[1]+=j
        return edge_new

    def matrix_loop(self, edge_list):
        # 4个邻接矩阵并列。
        assert len(edge_list) == 4
        A1, A2, A3, A4 = edge_list
        n = self.max_node_per_graph
        loop_edge_list = []
        loop_edge_list.append(A1)
        loop_edge_list.append(self.matrix_transfer(A2, n, 0))
        loop_edge_list.append(self.matrix_transfer(A3, 2*n, 0))
        loop_edge_list.append(self.matrix_transfer(A4, 3*n, 0))

        loop_edge_list.append(self.matrix_transfer(A4, 0, n))
        loop_edge_list.append(self.matrix_transfer(A1, n, n))
        loop_edge_list.append(self.matrix_transfer(A2, 2*n, n))
        loop_edge_list.append(self.matrix_transfer(A3, 3*n, n))

        loop_edge_list.append(self.matrix_transfer(A3, 0, 2*n))
        loop_edge_list.append(self.matrix_transfer(A4, n, 2*n))
        loop_edge_list.append(self.matrix_transfer(A1, 2*n, 2*n))
        loop_edge_list.append(self.matrix_transfer(A2, 3*n, 2*n))

        loop_edge_list.append(self.matrix_transfer(A2, 0, 3*n))
        loop_edge_list.append(self.matrix_transfer(A3, n, 3*n))
        loop_edge_list.append(self.matrix_transfer(A4, 2*n, 3*n))
        loop_edge_list.append(self.matrix_transfer(A1, 3*n, 3*n))

        return torch.cat(loop_edge_list, dim=1)

    def matrix_loop_new(self, edge_list,nodes_number):
        n = nodes_number
        edge_nums = len(edge_list)
        loop_edge_list = []
        edge_start = 0
        for j in range(edge_nums):
            cur_edge_start = edge_start    
            for i in range(edge_nums):
                cur_edge_start %= edge_nums
                loop_edge_list.append(self.matrix_transfer(edge_list[cur_edge_start], i*n, j*n))
                cur_edge_start += 1
            edge_start-=1
        return torch.cat(loop_edge_list, dim=1)


if __name__=="__main__":
    import torch
    tokens = torch.randint(low=0, high=69, size=(4,50, 19))
    edges= [torch.randint(low=0, high=50, size=(4,2,100)) for _ in range(4)]#[,torch.randint(low=0,high=9,size=10,),torch.randint(low=0,high=9,size=10,)]
    #edges1=[torch.randint(low=0, high=50, size=(2,100)) for _ in range(4)]
    model=Tensor_GCN(
            num_edge_types=4,
            in_features=19,
            out_features=40,
            embedding_out_features=40,
            embedding_num_classes=69,
            dropout=0.5,
            max_node_per_graph=50,
            device="cpu"
        )
    model.eval()
    
    new_edges=[torch.cat([es[j]+j*50 for j in range(4)],1) for es in edges]
    
    output=model(tokens,new_edges)
    
    for i in range(4):
        output2=model(tokens[i],[e[i] for e in edges])
        print(output.size(),output2.size())
        print(torch.sum(torch.abs(output[i]-output2[0])))
    